import argparse
import os
import ssl
import json
import re
from aiohttp import web
from yarl import URL

# ---------------- 工具函数 ----------------


def ensure_tokens_dir(tokens_dir: str):
    if not os.path.exists(tokens_dir):
        os.makedirs(tokens_dir)


def token_file_path(tokens_dir: str, token: str) -> str:
    return os.path.join(tokens_dir, f"{token}.json")


def load_token_data(tokens_dir: str, token: str) -> dict:
    path = token_file_path(tokens_dir, token)
    if os.path.exists(path):
        try:
            with open(path, "r") as f:
                return json.load(f)
        except Exception:
            return {}
    return {}


def save_token_data(tokens_dir: str, token: str, data: dict):
    path = token_file_path(tokens_dir, token)
    with open(path, "w") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


def is_valid_name(name: str) -> bool:
    return bool(re.match(r"^[a-zA-Z0-9_-]{1,64}$", name))


RESERVED_NAMES = {"login", "logout", "set", "delete", 'to', 'toggle'}

# ---------------- 登录相关 ----------------


async def login_page(request):
    return web.Response(text="""
        <html><body>
        <h2>Login</h2>
        <form method="post" action="/login">
            Token: <input type="password" name="token"/>
            <input type="submit" value="Login"/>
        </form>
        </body></html>
    """, content_type="text/html")


async def login_handler(request):
    data = await request.post()
    token = data.get("token", "")
    tokens_dir = request.app["tokens_dir"]
    if os.path.exists(token_file_path(tokens_dir, token)):
        resp = web.HTTPFound("/")
        resp.set_cookie("auth", token, httponly=True)
        return resp
    return web.Response(text="Invalid token", status=401)


async def logout_handler(request):
    resp = web.HTTPFound("/login")
    resp.del_cookie("auth")
    return resp


def get_cookie_token(request):
    return request.cookies.get("auth")


def check_token_header(request):
    auth = request.headers.get("Authorization", "")
    if auth.startswith("Bearer "):
        return auth.split(" ", 1)[1]
    return None

# ---------------- 页面操作 ----------------


async def list_handler(request):
    token = get_cookie_token(request)
    if not token:
        return web.HTTPFound("/login")

    query = request.query.get("q", "").strip().lower()
    sort_by = request.query.get("sort", "name")
    order = request.query.get("order", "asc")

    records = load_token_data(request.app['tokens_dir'], token)

    filtered = {
        name: record for name, record in records.items()
        if query in name.lower()
    }

    reverse = (order == "desc")
    if sort_by == "public":
        sorted_items = sorted(filtered.items(), key=lambda x: x[1].get(
            "public", False), reverse=reverse)
    else:
        sorted_items = sorted(
            filtered.items(), key=lambda x: x[0].lower(), reverse=reverse)

    html = [f"<html><body><h2>{token}的映射列表</h2>"]
    html.append('<p><a href="/logout">退出登录</a></p>')

    html.append(f"""
    <!-- 筛选表单 -->
    <form method="get" action="/" style="margin-bottom:16px;">
        <input type="text" name="q" placeholder="搜索 name..." value="{query}"/>
        <select name="sort" onchange="this.form.submit()">
            <option value="name" {"selected" if sort_by == "name" else ""}>按名称</option>
            <option value="public" {"selected" if sort_by == "public" else ""}>按公开状态</option>
        </select>
        <select name="order" onchange="this.form.submit()">
            <option value="asc" {"selected" if order == "asc" else ""}>升序</option>
            <option value="desc" {"selected" if order == "desc" else ""}>降序</option>
        </select>
    </form>

    <!-- 新增/更新记录表单 -->
    <form method="post" action="/set" style="padding:12px;border:1px solid #ddd;margin-bottom:16px;">
        <div style="font-weight:bold;margin-bottom:8px;">新增/更新记录</div>
        <label>名称 (a-zA-Z0-9_-): 
            <input type="text" name="name" required pattern="[a-zA-Z0-9_-]{1,64}" />
        </label>
        <br/>
        <label>URL: 
            <input type="url" name="url" required placeholder="https://example.com/page"/>
        </label>
        <br/>
        <label>
            <input type="checkbox" name="public"/> 公开
        </label>
        <br/>
        <input type="submit" value="保存"/>
    </form>
""")

    html.append('<ul>')
    for name, record in sorted_items:
        url = record.get("url", "#")
        public = record.get("public", False)
        toggle_label = "私有" if not public else "公开"

        html.append(f"""
            <li>
                {name}: 
                <a href="/{name}" target="_blank">{url}</a>
                [<a href="/delete/{name}">删除</a>] 
                [<a href="/toggle/{name}">{toggle_label}</a>] 
                <button type="button" onclick="enableEdit('{name}')">编辑</button>

                <form method="post" action="/set" id="form-{name}" style="display:none; margin-top:4px;">
                    <input type="hidden" name="name" value="{name}"/>
                    <input type="url" name="url" value="{url}" required onkeydown="if(event.key==='Enter') this.form.submit();"/>
                    <label><input type="checkbox" name="public" {"checked" if public else ""}/> 公开</label>
                </form>
            </li>
        """)

    html.append('</ul>')
    html.append('</body></html>')
    html.append("""
<script>
function enableEdit(name) {
    document.getElementById('form-' + name).style.display = 'block';
}
</script>
""")

    return web.Response(text="".join(html), content_type="text/html")


async def delete_handler(request):
    token = get_cookie_token(request)
    if not token:
        return web.HTTPFound("/login")

    token_data = load_token_data(request.app['tokens_dir'], token)
    name = request.match_info.get("name")
    if name in token_data:
        del token_data[name]
        save_token_data(request.app["tokens_dir"], token, token_data)
        return web.HTTPFound("/")
    return web.Response(text="not found", status=404)

# ---------------- API 接口 ----------------
async def set_handler(request):
    # 优先取 Bearer Token，没有就用 Cookie
    token = check_token_header(request)
    if not token:
        token = get_cookie_token(request)

    if not token:
        return web.json_response({"error": "unauthorized"}, status=401)

    content_type = request.headers.get("Content-Type", "")
    data = {}
    try:
        if content_type.startswith("application/json"):
            data = await request.json()
        else:
            post = await request.post()
            data = {
                "name": post.get("name", ""),
                "url": post.get("url", ""),
                "public": post.get("public", "") in ("1", "true", "on", "yes")
            }
    except Exception:
        return web.json_response({"error": "invalid payload"}, status=400)

    name = data.get("name", "").strip()
    url = data.get("url", "").strip()
    is_public = bool(data.get("public", False))

    if not name or not url:
        return web.json_response({"error": "name and url required"}, status=400)
    if not is_valid_name(name):
        return web.json_response({"error": "invalid name format"}, status=400)
    if name in RESERVED_NAMES:
        return web.json_response({"error": f"name '{name}' is reserved"}, status=400)

    token_data = load_token_data(request.app['tokens_dir'], token)
    token_data[name] = {"url": url, "public": is_public}
    save_token_data(request.app["tokens_dir"], token, token_data)

    # 如果是表单提交，直接跳回列表页
    if not content_type.startswith("application/json"):
        return web.HTTPFound("/")   # 关键修改点

    # API 调用才返回 JSON
    return web.json_response({"status": "ok", "data": {name: url, "public": is_public}})

# ---------------- 跳转接口 ----------------


async def redirect_handler(request):
    token = get_cookie_token(request)
    if not token:
        return web.Response(text="not logged in", status=401)

    token_data = load_token_data(request.app['tokens_dir'], token)
    name = request.match_info.get("name")
    record = token_data.get(name)
    if record:
        raise web.HTTPFound(record.get("url"))
    return web.Response(text="not found", status=404)


async def redirect_by_token_handler(request):
    token = request.match_info.get("token")
    name = request.match_info.get("name")
    extra = request.match_info.get("extra", "")  # 新增的路径部分
    token_data = load_token_data(request.app['tokens_dir'], token)

    record = token_data.get(name)
    if not record:
        return web.Response(text="not found", status=404)
    if not record.get("public"):
        return web.Response(text="not public", status=403)

    base_url = URL(record["url"])

    # 拼接额外路径
    if extra:
        base_url = base_url / extra

    # 合并 query 参数
    merged_url = base_url.with_query({**base_url.query, **request.query})

    raise web.HTTPFound(str(merged_url))


async def toggle_handler(request):
    token = get_cookie_token(request)
    if not token:
        return web.HTTPFound("/login")

    name = request.match_info.get("name")
    token_data = load_token_data(request.app["tokens_dir"], token)
    record = token_data.get(name)
    if not record:
        return web.Response(text="not found", status=404)

    record["public"] = not record.get("public", False)
    save_token_data(request.app["tokens_dir"], token, token_data)
    return web.HTTPFound("/")

# ---------------- 主程序 ----------------


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8080)
    parser.add_argument("--ssl-cert")
    parser.add_argument("--ssl-key")
    parser.add_argument("--tokens-dir", default="tokens",
                        help="存放 token.json 文件的目录")
    args = parser.parse_args()

    ensure_tokens_dir(args.tokens_dir)

    ssl_context = None
    if args.ssl_cert and args.ssl_key:
        ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        ssl_context.load_cert_chain(args.ssl_cert, args.ssl_key)

    app = web.Application()
    app["tokens_dir"] = args.tokens_dir

    # 路由注册
    app.router.add_get("/", list_handler)
    app.router.add_get("/login", login_page)
    app.router.add_post("/login", login_handler)
    app.router.add_get("/logout", logout_handler)
    app.router.add_get("/delete/{name}", delete_handler)
    app.router.add_post("/set", set_handler)
    app.router.add_get("/{name}", redirect_handler)  # 登录跳转
    app.router.add_get("/to/{token}/{name}", redirect_by_token_handler)  # 公开跳转
    app.router.add_get("/to/{token}/{name}/{extra:.*}", redirect_by_token_handler)
    app.router.add_get("/toggle/{name}", toggle_handler)

    web.run_app(app, host=args.host, port=args.port, ssl_context=ssl_context)


if __name__ == "__main__":
    main()
